package org.qiyu.framework.protect.limit.service.impl;

import org.qiyu.framework.protect.limit.service.CounterLimit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 滑动窗口限流
 *
 * @Author linhao
 * @Date created in 4:27 下午 2022/3/1
 */
public class SlidingWindowCounterLimit extends CounterLimit {

    private static Logger logger = LoggerFactory.getLogger(SlidingWindowCounterLimit.class);

    private AtomicInteger[] gridDistribution;

    private volatile int currentIndex = 0;

    /**
     * 当前时间之前的滑动窗口计数
     */
    private int preTotalCount;
    /**
     * 格子数
     */
    private int gridNumber;
    /**
     * 是否正在执行状态重置
     */
    private volatile boolean resetting = false;


    public SlidingWindowCounterLimit(int limitCount, int limitTime) {
        this(limitCount, limitCount, limitTime, TimeUnit.SECONDS);
    }

    public SlidingWindowCounterLimit(int gridNumber, int limitCount, int limitTime, TimeUnit timeUnit) {
        this.gridNumber = gridNumber;
        this.limitCount = limitCount;
        this.limitTime = limitTime;
        this.timeUnit = timeUnit;
        gridDistribution = new AtomicInteger[gridNumber];
        for (int i = 0; i < gridNumber; i++) {
            gridDistribution[i] = new AtomicInteger(0);
        }
//        new Thread(new CounterResetThread()).start();
    }

    @Override
    public boolean tryCount() {
        while (true) {
            if (isLimited) {
                logger.info("已经达到请求阈值，限流：{}", LocalDateTime.now().toString());
                return false;
            } else {
                //如果处于重置状态，则不要参与计算
                if(resetting){
                    logger.info("当前处于重置状态：{}", LocalDateTime.now().toString());
                    continue;
                }
                int currentGridCount = gridDistribution[currentIndex].get();
                //确保任意limitTime内都只有limitCount个请求
                if (preTotalCount + currentGridCount == limitCount) {
                    logger.info("当前请求正好触发请求阈值，限流：{}", LocalDateTime.now().toString());
                    isLimited = true;
                    return false;
                }
                //如果是重置状态则会进入下一个循环
                if (!resetting && gridDistribution[currentIndex].compareAndSet(currentGridCount, currentGridCount + 1)) {
                    return true;
                }
            }
        }
    }

    class CounterResetThread implements Runnable {

        @Override
        public void run() {
            resetting = false;
            int j = 0;
            while (true) {
                try {
                    timeUnit.sleep(1);
                    //当前时间格内的请求量
                    if (currentIndex >= gridNumber) {
                        resetting = true;
                        //重置数组
                        for (int i = 0; i < gridDistribution.length; i++) {
                            int temp = gridDistribution[i].get();
                            gridDistribution[i].compareAndSet(temp, 0);
                        }
                        j++;
                        preTotalCount = 0;
                        currentIndex = 0;
                        resetting = false;
                        isLimited = false;
                    }
                    int currentGridCount = gridDistribution[currentIndex].get();
                    preTotalCount = preTotalCount + currentGridCount;
                    logger.info("当前时间格内的请求量:{},第{}个时间单元内已请求量:{}", currentGridCount, j, preTotalCount);
                    currentIndex++;
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
